from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
import joblib
import numpy as np
from data_augment import augment_image

# 加载 MNIST 数据集
mnist = fetch_openml('mnist_784', parser="auto")

# 数据和标签
X, y = mnist['data'].values, mnist['target'].values

# 将数据重塑为28x28图像以便进行数据增强
X = X.reshape(-1, 28, 28)

# 分割数据集为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 数据归一化：只在训练集上拟合scaler
scaler = StandardScaler()
X_train_flat = X_train.reshape(X_train.shape[0], -1)  # 将训练集展平以便缩放
X_train_scaled = scaler.fit_transform(X_train_flat)

# 使用相同的scaler对测试集进行转换
X_test_flat = X_test.reshape(X_test.shape[0], -1)  # 将测试集展平以便缩放
X_test_scaled = scaler.transform(X_test_flat)


# 使用数据增强生成额外的训练样本
augmented_X = []
augmented_y = []

for i in range(len(X_train)):
    # 增强后的图像
    augmented_image = augment_image(X_train[i])
    augmented_X.append(augmented_image)  # 增强后的图像
    augmented_y.append(y_train[i])  # 原始标签

# 确保 augmented_X 中的所有元素都具有相同的形状 (28, 28)
augmented_X = np.array(augmented_X)  # 将增强图像转换为NumPy数组
assert augmented_X.shape[1:] == (28, 28), f"图像形状不一致: {augmented_X.shape}"

# 将原始训练集和增强的训练集结合
X_train_augmented = np.concatenate([X_train, augmented_X], axis=0)
y_train_augmented = np.concatenate([y_train, np.array(augmented_y)], axis=0)

# 重新展平增强后的训练集
X_train_augmented_flat = X_train_augmented.reshape(X_train_augmented.shape[0], -1)
X_train_augmented_scaled = scaler.transform(X_train_augmented_flat)

# 初始化KNN分类器
knn = KNeighborsClassifier(n_neighbors=3)

# 训练模型
knn.fit(X_train_augmented_scaled, y_train_augmented)

# 预测
y_pred = knn.predict(X_test_scaled)

# 计算准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"模型的准确率: {accuracy:.2f}")

# 保存模型和scaler
joblib.dump(knn, 'models/knn_mnist_model.pkl')
joblib.dump(scaler, 'models/scaler.pkl')
